<template>
  <d2-container>
    <template slot="header">猫狗识别模型训练</template>
    <div class="model_container">
      <el-row :gutter="20">
        <el-col :span="12">
          <div class="model_container_left">
            <el-card class="model_card" shadow="always">
              <div slot="header">
                模型训练数据资源池
              </div>

              <image-list/>
            </el-card>
          </div>
        </el-col>
        <el-col :span="12">
          <el-card class="model_card" shadow="always">
            <div class="train_logs flex_align" slot="header">
              <div>
                训练情况 <span class="tip_message">tips:模型训练时,浏览器计算压力大,页面将稍有卡顿</span>
              </div>
            </div>

            <div class="train_box">
              <div class="train_options">

                <div class="train_options_progress">
                  <el-progress :percentage="progressPercentage" :status="progressStatus"></el-progress>
                </div>
                <el-button type="primary" @click="initializeTraining" :loading="isTraining">
                  {{ isTraining ? '训练中' : model ? '重新训练' : '开始训练' }}
                </el-button>
                <el-button type="primary" @click="testDialogVisible = true" :disabled="!model">
                  模型测试
                </el-button>
              </div>
              <div class="train_box_logs" max-height="300px">
                <el-table
                  :data="trainingLogs"
                  style="width: 100%">
                  <el-table-column
                    prop="epoch"
                    label="轮数"
                  >
                  </el-table-column>
                  <el-table-column
                    prop="loss"
                    label="损失"
                  >
                  </el-table-column>
                  <el-table-column
                    prop="accuracy"
                    label="准确率">
                    <template slot-scope="{row}">
                      {{ (row.accuracy * 100).toFixed(1) }}%
                    </template>
                  </el-table-column>
                </el-table>
              </div>
            </div>
          </el-card>

        </el-col>
      </el-row>
    </div>

    <test-model-dialog
      v-if="testDialogVisible"
      v-bind="{ model,
                tf: tfComputed,
              }"
      @cancel="testDialogVisible = false"
    />
  </d2-container>
</template>

<script>
import * as tf from '@tensorflow/tfjs'
import { preprocessImage } from '@/views/demo/page1/utils/formatImages'
import TestModelDialog from '@/views/demo/page1/dialog/test-model-dialog.vue'
import ImageList from '@/views/demo/page1/components/image-list.vue'

export default {
  name: 'page1',
  components: {
    TestModelDialog,
    ImageList
  },
  computed: {
    tfComputed () {
      return tf
    }
  },
  data () {
    return {

      progressPercentage: 0, // 进度百分比
      progressStatus: null, // 进度条状态
      progressInterval: null, // 定时器
      isProgressComplete: false, // 新增：控制进度条是否完成
      testDialogVisible: false, //  测试弹窗
      model: null, // 模型
      testSamples: null, //
      selectedTestIndex: -1, // 选中的测试样本索引
      prediction: null, // 预测结果
      isTraining: false, // 是否正在训练
      trainingLogs: [] // 训练日志
    }
  },
  mounted () {
  },
  methods: {

    startProgress () {
      this.progressPercentage = 0
      this.progressStatus = null
      this.isProgressComplete = false
      this.progressInterval = setInterval(() => {
        if (this.progressPercentage < 99) {
          this.progressPercentage += 1
        } else {
          clearInterval(this.progressInterval)
        }
      }, 500)
    },
    // 完成进度条
    completeProgress () {
      clearInterval(this.progressInterval)
      this.progressPercentage = 100
      this.progressStatus = 'success'
      this.isProgressComplete = true
    },

    // 加载数据集
    async loadDataset () {
      const images = []
      const categories = ['cat', 'dog']
      for (const category of categories) {
        for (let i = 1000; i < 1500; i++) {
          try {
            const imgData = await this.loadImage(category, i)
            images.push(imgData)
          } catch (error) {
            console.error(
              `加载${
                category === 'cat' ? '猫' : '狗'
              }图片失败: ${category}.${i}.jpg`,
              error
            )
          }
        }
      }
      return images
    },

    async loadImage (category, i) {
      const imgPath = require(`@/views/demo/page1/assets/image/${category}/${category}.${i}.jpg`)
      const img = new Image()
      img.src = imgPath.default || imgPath
      await new Promise((resolve, reject) => {
        img.onload = () => resolve(img)
        img.onerror = (err) => {
          console.error(`图片加载失败: ${imgPath}`)
          reject(err)
        }
      })
      return {
        path: imgPath,
        element: img,
        tensor: tf.browser.fromPixels(preprocessImage(img)).div(255), // 归一化
        label: category === 'cat' ? 0 : 1
      }
    },
    // 创建卷积神经网络模型

    createCNNModel () {
      const model = tf.sequential({
        layers: [
          // 最大池化层：降低特征图尺寸，增强特征鲁棒性
          tf.layers.maxPooling2d({
            inputShape: [128, 128, 3], // 输入形状 [高度, 宽度, 通道数]
            poolSize: 2, // 池化窗口尺寸 2x2
            strides: 2 // 滑动步长：每次移动 n 像素，使输出尺寸减小到原先的 1/n
          }),

          // 卷积层：用于提取图像局部特征
          tf.layers.conv2d({
            filters: 32, // 卷积核数量，决定输出特征图的深度
            kernelSize: 3, // 卷积核尺寸 3x3
            activation: 'relu', // 激活函数：修正线性单元，解决梯度消失问题
            padding: 'same' // 边缘填充方式：保持输出尺寸与输入相同
          }),
          // 展平层：将多维特征图转换为一维向量
          tf.layers.flatten(),

          // 全连接层（输出层）：进行最终分类
          tf.layers.dense({
            units: 2, // 输出单元数：对应猫/狗两个类别
            activation: 'softmax' // 激活函数：将输出转换为概率分布
          })
        ]
      })

      // 编译模型：配置训练参数
      model.compile({
        optimizer: 'adam',
        loss: 'categoricalCrossentropy',
        metrics: ['accuracy']
      })

      console.log('模型架构:')
      model.summary()

      return model
    },

    // 准备训练数据
    prepareTrainingData (images) {
      const tensors = images.map((img) => img.tensor)
      return {
        xData: tf.stack(tensors),
        yData: tf.oneHot(
          tf.tensor1d(
            images.map((img) => img.label),
            'int32'
          ),
          2
        )
      }
    },
    // 训练模型
    async trainModel (model, xData, yData) {
      this.trainingLogs = [] // 清空之前的训练日志
      await model.fit(xData, yData, {
        epochs: 10, // 训练轮数
        batchSize: 4,
        validationSplit: 0.4,
        callbacks: {
          onEpochEnd: (epoch, logs) => {
            const log = {
              epoch: epoch + 1,
              loss: Number(logs.loss.toFixed(4)),
              accuracy: Number(logs.acc.toFixed(4))
            }
            this.trainingLogs.push(log)
          }
        }
      })
    },

    //  预测函数
    predict (imgData) {
      if (!this.model) return
      const input = imgData.tensor.expandDims(0) // 添加批次维度
      const prediction = this.model.predict(input)
      const probabilities = Array.from(prediction.dataSync())
      this.prediction = {
        label: probabilities[0] > probabilities[1] ? '猫' : '狗',
        confidence: Math.max(...probabilities) * 100
      }
      tf.dispose([input, prediction])
    },

    // 初始化训练过程
    async initializeTraining () {
      this.startProgress()
      try {
        this.isTraining = true
        this.trainingLogs = []
        const allImages = await this.loadDataset()

        // 分割训练集和测试集
        const shuffledImages = [...allImages]
        tf.util.shuffle(shuffledImages)
        const testSet = shuffledImages.slice(-10) // 取最后4个作为测试集
        const trainSet = shuffledImages.slice(0, -10)
        this.testSamples = testSet
        const {
          xData,
          yData
        } = this.prepareTrainingData(trainSet)
        const newModel = this.createCNNModel()
        await this.trainModel(newModel, xData, yData)
        this.model = newModel
        this.$message.success('模型训练完成')
        // 训练完成后，默认选择第一个测试样本
        if (testSet.length > 0) {
          this.selectedTestIndex = 0
          this.predict(testSet[0])
        }
      } catch (error) {
        console.error('训练初始化失败:', error)
        this.$message.error('训练初始化失败:', error)
      } finally {
        this.isTraining = false
        this.completeProgress()
      }
    },
    // 选择测试样本
    handleTestSampleChange (index) {
      this.selectedTestIndex = index
      this.predict(this.testSamples[index])
    }

  },
  beforeDestroy () {
    // 组件销毁前清除定时器
    clearInterval(this.progressInterval)
  }
}
</script>

<style scoped lang="scss">
.flex_align {
  display: flex;
  align-items: center;
}

.train_logs {
  justify-content: space-between;
}

.model_card {
  height: 60vh;

  .el-card__body {
  }

}

.image-gallery {
  display: flex;
  flex-wrap: wrap;
  gap: 10px;
}

.image-item {
  width: 150px;
  text-align: center;
}

.gallery-image {
  width: 100%;
  height: 150px;
  object-fit: cover;
  border-radius: 4px;
}

.image-label {
  margin-top: 5px;
  font-size: 14px;
  color: #666;
}

.train_box_logs {
  margin-top: 20px;
}

.tip_message {
  font-size: 12px;
  color: #909399;
}

.train_options {
  display: flex;
  align-items: center;
  justify-content: flex-end;

  .train_options_progress {
    width: 80%;
  }
}
</style>
